#pragma once

#include <array>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <exception>
#include <memory>
#include <new>
#include <stdexcept>
#include <type_traits>
#include <utility>
#include <vector>

namespace Modern {

/**
 * @brief 改进版固定容量内存池，实现显式生命周期管理、完美转发等。
 */
template <typename T>
class MemPool final {
private:
  struct Block {
    alignas(T) std::array<std::byte, sizeof(T)> storage{};
    bool is_free{true};

    [[nodiscard]] T* raw_address() noexcept {
      return reinterpret_cast<T*>(storage.data());
    }

    [[nodiscard]] const T* raw_address() const noexcept {
      return reinterpret_cast<const T*>(storage.data());
    }

    [[nodiscard]] T* object_ptr() noexcept {
      return std::launder(raw_address());
    }

    [[nodiscard]] const T* object_ptr() const noexcept {
      return std::launder(raw_address());
    }
  };

  static_assert(std::is_standard_layout_v<Block>,
                "Block 必须是标准布局。");
  static_assert(offsetof(Block, storage) == 0,
                "storage 必须是 Block 的首个成员。");

public:
  explicit MemPool(std::size_t capacity)
      : blocks_(capacity) {
    if (capacity == 0) {
      throw std::invalid_argument("MemPool 的容量必须大于零。");
    }
    free_list_.reserve(capacity);
    for (std::size_t idx = capacity; idx-- > 0;) {
      free_list_.push_back(idx);
    }
  }

  ~MemPool() {
    for (auto& block : blocks_) {
      if (!block.is_free) {
        std::destroy_at(block.object_ptr());
      }
    }
  }

  MemPool() = delete;
  MemPool(const MemPool&) = delete;
  MemPool& operator=(const MemPool&) = delete;
  MemPool(MemPool&&) = delete;
  MemPool& operator=(MemPool&&) = delete;

  /**
   * @brief 分配一个对象，支持完美转发构造参数。
   * @throws std::bad_alloc 当内存池已满。
   */
  template <typename... Args>
  requires std::constructible_from<T, Args...>
  [[nodiscard]] T* allocate(Args&&... args) {
    if (free_list_.empty()) {
      throw std::bad_alloc();
    }

    const auto index = free_list_.back();
    free_list_.pop_back();
    auto& block = blocks_[index];
    assert(block.is_free && "分配时块应为空闲。");

    T* address = block.raw_address();

    try {
      std::construct_at(address, std::forward<Args>(args)...);
      block.is_free = false;
      ++live_objects_;
      return address;
    } catch (...) {
      free_list_.push_back(index);
      throw;
    }
  }

  /**
   * @brief 回收对象：调用析构函数并把块放回空闲栈。
   */
  void deallocate(T* ptr) noexcept {
    if (!ptr) {
      return;
    }
    if (!owns(ptr)) {
      assert(false && "指针不属于当前内存池。");
      std::terminate();
    }

    const auto index = index_from_pointer(ptr);
    auto& block      = blocks_[index];
    if (block.is_free) {
      assert(false && "检测到重复释放。");
      std::terminate();
    }

    std::destroy_at(ptr);
    block.is_free = true;
    free_list_.push_back(index);
    --live_objects_;
  }

  [[nodiscard]] bool owns(const T* ptr) const noexcept {
    if (!ptr || blocks_.empty()) {
      return false;
    }

    const auto base =
        reinterpret_cast<std::uintptr_t>(blocks_.data());
    const auto value =
        reinterpret_cast<std::uintptr_t>(ptr);
    const auto total =
        static_cast<std::uintptr_t>(blocks_.size()) * sizeof(Block);

    if (value < base || value >= base + total) {
      return false;
    }

    const auto offset = value - base;
    return (offset % sizeof(Block)) == 0;
  }

  [[nodiscard]] std::size_t capacity() const noexcept {
    return blocks_.size();
  }

  [[nodiscard]] std::size_t size() const noexcept {
    return live_objects_;
  }

  [[nodiscard]] bool empty() const noexcept {
    return live_objects_ == 0U;
  }

private:
  [[nodiscard]] std::size_t index_from_pointer(const T* ptr) const noexcept {
    const auto base =
        reinterpret_cast<std::uintptr_t>(blocks_.data());
    const auto value =
        reinterpret_cast<std::uintptr_t>(ptr);
    const auto offset = value - base;
    return static_cast<std::size_t>(offset / sizeof(Block));
  }

  std::vector<Block> blocks_;
  std::vector<std::size_t> free_list_;
  std::size_t live_objects_{0};
};

} // namespace Modern